/*
    Copyright 2006-2012 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Declaration of the full_sed_grid class.

#ifndef __full_sed_grid__
#define __full_sed_grid__

#include "polychromatic_grid.h"
#include "optical.h"
#include "blitz-fits.h"
#include "mcrx-units.h"
#include "fits-utilities.h"
#include "CCfits/CCfits"
#include "mcrx-debug.h"
#include "counter.h"
#include "constants.h"
#include "units.h"
#include "mpi_util.h"

namespace mcrx {
  template<typename> class full_sed_grid;
  // forward decl
  class density_generator;
  template<typename T> class arepo_grid;

  // typedefs for commonly used full_sed_grid types
  typedef full_sed_grid<
    adaptive_grid<
      cell_data<
	typename emitter_types<adaptive_grid, 
			       polychromatic_policy, local_random>::T_emitter,
	absorber<array_1> > > > T_full_sed_adaptive_grid;
#ifdef WITH_AREPO
  typedef full_sed_grid<
    arepo_grid<
      cell_data<
	typename emitter_types<arepo_grid, 
			       polychromatic_policy, local_random>::T_emitter,
	absorber<array_1> > > > T_full_sed_arepo_grid;
#endif
}

// forward decl
class binofstream;
class binifstream;

/** A polychromatic grid that knows about units, used for the actual
    mcrx runs.
 */
template<typename grid_type>
class mcrx::full_sed_grid : public mcrx::polychromatic_grid<grid_type>
{
  std::pair<CCfits::ExtHDU*, CCfits::ExtHDU*>
  create_hdus(CCfits::FITS& file, const blitz::TinyVector<int,2>& size,
	      const std::string& i_name, const std::string& d_name,
	      const string& luminosityunit, bool do_compress) const;

public:
  typedef grid_type T_grid;
  typedef polychromatic_grid<T_grid> T_base;
  typedef typename T_base::T_grid_impl T_grid_impl;
  typedef typename T_base::T_data T_data;
  typedef typename T_base::T_lambda T_lambda;
  typedef typename T_base::T_cell T_cell;
  typedef typename T_base::T_cell_tracker T_cell_tracker;
  typedef typename T_base::iterator iterator;
  typedef typename T_base::const_iterator const_iterator;

  /// Units
  T_unit_map units_;
  
  /** Constructor loads grid structure and cell data from a FITS file.
      The grid structure is loaded from a GRIDSTRUCTURE HDU by the base
      class constructor, and then the grid-cell SED's are loaded from a
      binary FITS table. */
  full_sed_grid (boost::shared_ptr<T_grid_impl> g,
		 CCfits::ExtHDU& structure_hdu,
		 CCfits::ExtHDU& data_hdu,
		 const density_generator& gen,
		 int n_lambda);

  /** This constructer loads cell data from the Arepo data
      structures. The grid structure is handled by the supplied grid
      object, all that we do here is use the SPH data to populate our
      data. (This can't be done by the polychromatic_grid constructor
      because we need to convert the Arepo units and only this class
      knows about units.) Note that no cell data is read from the
      sfrhist output file when using Arepo. */
  full_sed_grid (boost::shared_ptr<T_grid_impl> g,
		 const density_generator& gen,
		 int n_lambda);

  /** This constructor is used when a grid has been built using the
      factory. */
  full_sed_grid (boost::shared_ptr<T_grid_impl> g, const T_unit_map& units);

  template<typename dust_model>
  void write_intensity (CCfits::FITS&, 
			const std::string&, const std::string&,
			T_float normalization, 
			dust_model&, 
			const string& luminosityunit,
			bool do_compress);
  void load_intensity (const CCfits::ExtHDU&);
  void write_dump (binofstream& file) const;
  void load_dump (binifstream& file);

  /** Returns the conversion factor from internal area units (usually
      kpc) to SI (m^2). This is used when converting the intensity
      array to physical units.  The convert function will throw if
      there is no length unit, which is ok, since without units we
      can't make this conversion. */
  T_float area_factor() const {
    const T_float to_m = units::convert(units_.get("length"), "m");
    return to_m*to_m;
  };

};


/** Writes the radiation intensity and energy deposition in the grid
    cells as FITS tile-compressed 2D images. See the mcrx-readme for
    documentation of the file format. i_name and d_name are the name
    of the HDUs to which the intensity and energy deposition arrays
    are written. powerunit is the unit of whatever is in the grid,
    which this class has no knowledge of. */
template<typename T>
template<typename dust_model>
void mcrx::full_sed_grid<T>::write_intensity (CCfits::FITS& file,
					      const std::string& i_name,
					      const std::string& d_name,
					      T_float normalization, 
					      dust_model& dm,
					      const string& luminosityunit,
					      bool do_compress)
{
  using namespace blitz;
  using namespace CCfits;

  if(this->intensities_.size()==0) {
    cout << "Can't save radiative intensity, not saved." << endl;
    return;
  }

  std::cout << "Calculating radiative intensity and energy deposition" << std::endl;

  //  We will collect all the data first, and then write. For
  //  efficiency's sake, we collect the data straight into
  //  column-major arrays so they can be written directly to the FITS
  //  file

  // Indices are (scatterer, lambda).

  const int n_lambda=this->begin()->data()->get_absorber().intensity().size();

  Array<float, 2> intensity(shape(this->n_cells(), n_lambda),
			    ColumnMajorArray<2>());
  Array<float, 2> deposition(intensity.shape(),
			     ColumnMajorArray<2>());
  
  int ii=0;
  counter ct(100);

  // conversion is hard-coded to use m in output file regardless of
  // what internal units are
  const T_float areafactor = area_factor();

  for (const_iterator c = this->begin (); c != this->end (); ++c, ++ii, ct++) {
    /*
    // mean intensity, determined as sum(I*dl)/(4*pi*V)

    // This is absorbed energy based on our independent estimate of J,
    // which is less noisy than the direct Monte Carlo estimate of the
    // energy itself...
    // Actually, dep = 4*pi*V*ao*J = 4*pi*V*ao*int/(4*pi*V) = ao*int,
    // which saves us from any unit conversion. dep will automatically have
    // units of (W/m)/kpc^3.
    */

    ASSERT_ALL(c->data()->get_absorber().intensity()==this->intensities_(ii,Range::all()));
    ASSERT_ALL(c->data()->get_absorber().densities()==this->rho_(ii,Range::all()));

    intensity(ii, Range::all()) = 
      log10(this->intensities_(ii, Range::all())*normalization/
	    (4*constants::pi*c.volume()*areafactor) + 
	     blitz::tiny(T_float()));
    deposition(ii, Range::all()) = 
      log10(dm.total_abs_length_to_absorption(this->rho_(ii,Range::all())) *
	    this->intensities_(ii, Range::all())*normalization + 
	    blitz::tiny(T_float()));
  }

  ASSERT_ALL(intensity==intensity);
  ASSERT_ALL(deposition==deposition);
  DEBUG(1,cout << "Intensity min/max: " << min(intensity) << ' ' << max(intensity) << endl;);
  
  std::pair<ExtHDU*, ExtHDU*> hdus;

  if(this->shared_domain()) {
    // shared domain. coadd arrays.
    mpi_sum_arrays(intensity);
    mpi_sum_arrays(deposition);
    if(is_mpi_master()) {
      std::cout << "Writing intensities and energy deposition" << std::endl;
      hdus=create_hdus(file, intensity.shape(), i_name, d_name, 
		       luminosityunit, do_compress);
      write(*hdus.first, intensity);
      hdus.first->writeChecksum();
      write(*hdus.second, deposition);
      hdus.second->writeChecksum();
    }
  }
  else {    // split domain. Need to write consecutively
    {
      Array<float, 2> out(mpi_stack_arrays(intensity, firstDim));
      if(is_mpi_master()) {
	std::cout << "Writing intensities and energy deposition" << std::endl;
	hdus=create_hdus(file, out.shape(), i_name, d_name, 
			 luminosityunit, do_compress);
	write(*hdus.first, out);
	hdus.first->writeChecksum();
      }
    }
    {
      Array<float, 2> out(mpi_stack_arrays(deposition, firstDim));
      if(is_mpi_master()) {
	write(*hdus.second, out);
	hdus.second->writeChecksum();
      }
    }
  }

  // free arrays
  intensity.free();
  deposition.free();

  if(is_mpi_master()) {
    hdus.first->addKey("normalization", normalization,
		       "Image normalization (internal usage)");
    hdus.first->addKey("area_factor", areafactor,
		       "Area unit conversion (internal usage)");
    
    hdus.second->addKey("normalization", normalization,
			"Image normalization (internal usage)");
    hdus.second->addKey("area_factor", areafactor,
			"Area unit conversion (internal usage)");
    const string depunit = luminosityunit+"/"+units_.get("length")+"^3";
    hdus.second->addKey("imunit", depunit,
			"Energy deposition rate in units of "+depunit);
  }
}


#endif

  
